import base64
import hashlib
import json
import os
import sys
import time
import traceback
import uuid
from io import BytesIO
from typing import Optional, Tuple

import cv2
import numpy as np
import math
import requests
from PIL import Image, ImageFile
from dotenv import load_dotenv
from openai import OpenAI
from volcenginesdkarkruntime import Ark
from volcenginesdkarkruntime._exceptions import ArkAPIError

from utils import sr_s3, sr_sqs, sr_utils, sqs_video
from utils.sr_s3 import download_single_video
from utils.sr_utils import image_path_to_data_url, video_compress_image_size_path
from utils.status_codes import StatusCodes
from volcengine.visual.VisualService import VisualService

load_dotenv(verbose=True, override=True)
api_key = os.getenv('DB_API_KEY')

secret_id = os.getenv('DB_SECRET_ID', None)
secret_key = os.getenv('DB_SECRET_KEY', None)

# 第三方 sdk
visual_service = VisualService()
# call below method if you don't set ak and sk in $HOME/.volc/config
# visual_service.set_ak(secret_id)
# visual_service.set_sk(secret_key)
visual_service.set_ak('AKLTNjdmYTI2MjRiMWQwNDNhNjliNDU0NzAyNzZkNjM3MDM')
visual_service.set_sk('WVRsa016WTJZalV6TWpRNU5HVmhZbUZrTjJReE5EZGhOV1l5WldJMFl6RQ==')
visual_service.set_connection_timeout(300)
visual_service.set_socket_timeout(300)

# 配置与初始化
ImageFile.LOAD_TRUNCATED_IMAGES = True  # 处理截断图片
app_env = os.getenv('APP_ENV', 'prod')

# 队列初始化
if app_env == 'dev':
    TASK_QUEUE = sqs_video.videoSQS.get_queue_by_name(QueueName='AIHubVideoTasks_Test')
else:
    TASK_QUEUE = sqs_video.videoSQS.get_queue_by_name(QueueName='AIHubVideoTasks')

# 状态码常量
SUCCESS = 200
TASK_FAILED = 2001
TIMEOUT = 2002
PARAMETER_ERROR = 2003
CLIENT_INIT_ERROR = 2004
IMAGE_PROCESS_ERROR = 2005


# --------------------------
# 通用工具函数
# --------------------------
def delete_temp_files(*file_paths):
    """删除临时文件（忽略不存在的文件）"""
    for path in file_paths:
        if path and os.path.exists(path):
            try:
                os.remove(path)
                print(f"已清理临时文件: {path}")
            except Exception as e:
                print(f"清理临时文件失败 {path}: {str(e)}")


def download_single_image(img_url, storage_service, bucket_name, save_dir="./upload"):
    """下载单张图片到本地并返回路径"""
    try:
        os.makedirs(save_dir, exist_ok=True)
        name = str(uuid.uuid4())
        ext = img_url[img_url.rfind('.'):] if '.' in img_url else '.jpg'
        save_path = os.path.join(save_dir, f"{name}{ext}")

        # 调用S3工具下载图片
        success, err = sr_s3.download_file(
            img_url, save_path, serviceName=storage_service, bucketName=bucket_name
        )
        if not success:
            return None, f"下载失败: {err}"

        # 压缩图片（确保符合模型要求）
        success, new_path = sr_utils.video_compress_image_size_path(save_path, 25, (300, 6000), (300, 6000))
        if not success:
            delete_temp_files(save_path)
            return None, "图片压缩失败" + save_path

        return new_path, None
    except Exception as e:
        return None, f"处理图片异常: {str(e)}"


def upload_video_to_s3(video_path, s3_storage, storage_service, bucket_name):
    """上传视频到R2并返回R2路径"""
    if not video_path or not os.path.exists(video_path):
        return None, "视频文件不存在"

    try:
        # 调用S3工具上传视频
        s3_key = sr_s3.upload_login_or_visitor(
            video_path, s3_storage, serviceName=storage_service, bucketName=bucket_name
        )
        return s3_key, None
    except Exception as e:
        return None, f"上传R2失败: {str(e)}"




def send_result_to_queue(queue_name, body, status, msg, results=None, start_time=None, local_time=None, local_video=None, ext="mp4"):
    """发送结果到指定队列"""
    if local_video is None:
        local_video = {
            "width": 0,
            "height": 0,
            "fps": 0,
            "tokens": "0",
        }
    try:
        results = results or {}
        msg_body = {
            'QueueResult': queue_name,
            'status': status,
            'taskId': body.get('taskId'),
            'taskType': body.get('taskType'),
            'appId': body.get('appId'),
            'machineId': 1,
            'width': local_video['width'],
            'height': local_video['height'],
            'tokens': local_video['tokens'],
            'fps': local_video['fps'],
            'results': results,
            'msgFetchEndTime': start_time,
            'start_time': start_time,
            'end_time': sr_utils.getTimes(local_time),
            'ext': ext,
            'msg': msg
        }
        print(msg_body)
        queue = sqs_video.videoSQS.get_queue_by_name(QueueName=queue_name)
        queue.send_message(MessageBody=json.dumps(msg_body))
        print(f"已发送结果到队列 {queue_name}: {status}")
    except Exception as e:
        print(f"发送队列消息失败: {str(e)}")


# --------------------------
# 视频生成核心逻辑
# --------------------------
def viduGenerator_doubao(body, local_time=None):
    """统一处理视频生成任务（入口函数）"""

    task_id = body.get('taskId', 'unknown')
    storage_service = body.get('storageService', 's3')
    bucket_name = body.get('bucketName')
    s3_storage = body.get('storedPrefix')
    model_type = body.get('modelType')
    prompt = body.get('prompt')
    template = body.get('template')
    try:
        first_image_url = body.get('first_image')
        last_image_url = body.get('last_image')
        first_image_url = sr_s3.get_url(first_image_url, serviceName=storage_service, bucketName=bucket_name)
        if last_image_url != "":
            last_image_url = sr_s3.get_url(last_image_url, serviceName=storage_service, bucketName=bucket_name)
        else:
            last_image_url = None
        create_params = {
            "model_type": model_type,
            "prompt": prompt,
            "first_image": first_image_url,
            "last_image": last_image_url,
            "template": template,
            "save_dir": "./download"
        }
        local_video = create_and_save_video(**create_params)

        if not local_video['save_path'] or not os.path.exists(local_video['save_path']):
            msg_body = {'status': local_video['status'], "width": 0, "height": 0, "url": None, 'msg': local_video['msg'], 'code': local_video['code']}
            return msg_body

        # --------------------------
        # 3. 上传视频到S3
        # --------------------------
        print(f"开始上传视频到S3: {local_video['save_path']}")
        s3_video_key, err = upload_video_to_s3(
            local_video['save_path'], s3_storage, storage_service, bucket_name
        )
        if not s3_video_key:
            msg_body = {'status': 0, "width": 0, "height": 0, "url": None, 'msg': f"视频上传r2失败: {err}"}
            return msg_body

        # --------------------------
        # 4. 发送成功结果到队列
        # --------------------------
        msg_body = {
            'status': 1,
            'width': local_video['width'],
            'height': local_video['height'],
            'local_path': local_video['save_path'],
            'fps': local_video['fps'],
            'url': s3_video_key,
            "tokens": local_video['tokens'],
            "model_mode": "std",  # 模型类型标准模式
            'msg': "success",
            "code": ""
        }
        return msg_body

    except Exception as e:
        err_msg = f"任务处理异常: {str(e)}"
        print(f"任务{task_id}异常: {traceback.format_exc()}")
        msg_body = {'status': 0, "width": 0, "height": 0, "url": None, 'msg': err_msg, 'code': 6001}
        return msg_body

    # finally:
    # video_path = local_video.get('save_path') if isinstance(local_video, dict) else None
    # 确保临时文件被清理
    # delete_temp_files(first_image_path, last_image_path, video_path)


def create_video_task(
        model_type: int,
        first_image: Optional[str] = None,
        last_image: Optional[str] = None,
        template: Optional[str] = None,
):
    """
    创建视频生成任务并轮询结果（优化版）

    参数说明：
    - model_type: 模型类型 ("1": 文生视频, "2": 文+图生视频, "3": 文+首尾帧生视频)
    - prompt: 文本提示词
    - first_image: 首帧图片路径（可选）
    - last_image: 尾帧图片路径（仅首尾帧模式需要）
    - duration: 视频时长（秒）
    - fps: 帧率
    - resolution: 分辨率（如"720p"）
    - ratio: 宽高比
    - watermark: 是否带水印
    - seed: 随机种子（-1为随机）

    返回：(状态码, 任务ID, 消息, 消耗Tokens, 视频URL)
    """
    # --------------------------
    # 1. 参数校验（前置检查，减少无效调用）
    # --------------------------
    if model_type == 1:
        model = "i2v_bytedance_effects_v1"
    elif model_type == 2:
        model = "i2v_template_cv_v2",
    else:
        model = "i2v_bytedance_effects_v1"
    # --------------------------
    # 3. 构建请求内容（简化逻辑，减少重复）

    if last_image is not None:
        first_image =first_image + "|" + last_image
    form = {
        "req_key": model,
        "image_input": first_image,
        "template_id": template
        # ...
    }
    start_time = time.time()
    try:
        resp = visual_service.cv_sync2async_submit_task(form)
        task_id = resp['data']['task_id']
        # task_id = 14231194916069442428
        print(f"任务创建成功，ID: {task_id}")
    except Exception as e:
        error_json = sr_utils.handle_exception(e)
        print(error_json.get("code"))
        err_code = error_json.get('code', '未知错误码')
        message = error_json.get('message', '未知错误码')
        return err_code, '', f"任务失败: {message}", None, None
    #
    queue_form = {
        "req_key": model,
        "task_id": task_id
        # ...
    }
    # --------------------------
    # 5. 轮询任务状态（优化轮询逻辑）
    # --------------------------
    max_retries = 600  # 10分钟（60*5s）

    retry_interval = 5  # 初始间隔5秒
    for retry in range(max_retries):
        try:
            task_status = visual_service.cv_sync2async_get_result(queue_form)
            status = task_status['data']['status']
            # 初始化 code 为 None（默认值）
            error_code = None
            code = task_status['code']
            message = task_status['message']
            # 成功状态
            if status == "done" and code == 10000 and message == 'Success':
                elapsed = time.time() - start_time
                video_url = json.loads(task_status['data']['resp_data'])['video_url']
                print(f"任务{task_id}成功，耗时{elapsed:.2f}秒，video_url: {video_url}")
                return SUCCESS, task_id, f"任务成功，耗时{elapsed:.2f}秒", 0, video_url
            # 失败状态
            if status == "done" and code != 10000 and message == 'Success':
                error_msg = task_status['message'] if task_status['message'] else "未知错误"
                print(f"任务{task_id}失败: {error_msg}")
                # 仅当 error 存在时，才获取 code（避免 error 为 None 时的异常）
                if task_status['code'] is not None:
                    error_code = task_status['code']
                return error_code, task_id, f"任务失败: {error_msg}", error_code, None

            # 进行中（打印进度，动态调整间隔）
            if retry % 6 == 0:  # 每30秒打印一次状态
                print(f"任务{task_id}状态: {status}（{retry + 1}/{max_retries}）")
            time.sleep(retry_interval)
        except Exception as e:
            error_json = sr_utils.handle_exception(e)
            error_code = error_json.get('code')
            if error_json.get("code") != 10000:
                error_msg = error_json['message'] if error_json['message'] else "未知错误"
                return error_code, task_id, f"任务失败: {error_msg}", error_code, None
            print(f"轮询任务{task_id}异常: {str(e)}，将重试...")
            time.sleep(retry_interval)  # 轮询异常时仍重试

    # 超时处理
    print(f"任务{task_id}超时（超过10分钟）")
    return TIMEOUT, task_id, "任务超时（超过10分钟）", None, None


def get_video_info(video_path: str) -> dict:
    """获取视频的宽高、时长和帧率信息"""
    try:
        cap = cv2.VideoCapture(video_path)

        # 获取视频的基本信息
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        fps = cap.get(cv2.CAP_PROP_FPS)
        frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

        # 计算视频时长（秒）
        duration = frame_count / fps if fps > 0 else 0

        cap.release()

        return {
            'width': width,
            'height': height,
            'duration': duration,
            'fps': fps
        }
    except Exception as e:
        print(f"获取视频信息失败：{str(e)}")
        # 返回默认值或部分信息
        return {
            'width': None,
            'height': None,
            'duration': None,
            'fps': None
        }


def download_video(video_url: str, save_dir: str = "./videos", filename: Optional[str] = None) -> Optional[dict]:
    """
    从视频URL下载视频并保存到本地

    参数：
    - video_url: 视频下载URL
    - save_dir: 保存目录（默认 ./videos）
    - filename: 自定义文件名（如不指定则从URL提取）

    返回：
    - 本地保存路径（成功）或 None（失败）
    """
    try:
        # 创建保存目录（如不存在）
        os.makedirs(save_dir, exist_ok=True)

        # 生成文件名（默认从URL提取）
        if not filename:
            # 从URL尾部提取文件名（如无则用随机名）
            filename = video_url.split("/")[-1]
            if not filename.endswith((".mp4", ".mov", ".avi")):
                filename = f"video_{os.urandom(4).hex()}.mp4"  # 随机名+默认mp4格式

        save_path = os.path.join(save_dir, filename)

        # 下载视频（带超时和进度提示）
        print(f"开始下载视频到 {save_path}...")
        response = requests.get(video_url, stream=True, timeout=600)
        response.raise_for_status()  # 检查HTTP错误
        # 分块写入文件（适合大文件）
        with open(save_path, "wb") as f:
            for chunk in response.iter_content(chunk_size=1024 * 1024):  # 1MB/块
                if chunk:
                    f.write(chunk)
            # 获取视频宽高
        video_info = get_video_info(save_path)
        print(f"视频保存成功：{save_path}")

        return {
            'save_path': save_path,
            'width': video_info['width'],
            'height': video_info['height'],
            'duration': video_info['duration'],
            'fps': video_info['fps']
        }
        # return save_path
    except Exception as e:
        print(f"视频下载失败：{str(e)}")
        return None


def create_and_save_video(
        model_type: int,
        save_dir: str = "./download/videos",
        first_image: Optional[str] = None,
        template: Optional[str] = None,
        last_image: Optional[str] = None,

):
    """
    生成视频并保存到本地（整合生成+下载流程）
    返回本地视频路径（失败则返回None）
    """
    # 1. 调用视频生成函数
    code, task_id, msg, tokens, video_url = create_video_task(
        model_type=model_type,
        first_image=first_image,
        last_image=last_image,
        template=template
    )
    video_result = {
        "width": 0,
        "height": 0,
        "fps": 0,
        "status": 0,
        "tokens": 0,
        "save_path": None,
        "msg": msg,
        "code": 6001,
    }
    # 2. 检查生成结果
    if code != 200 or not video_url:
        print(f"code:{code},视频生成失败：{msg} , 任务ID：{task_id},error:{tokens}")
        if code == "50411":
            err_code = StatusCodes.INPUT_IMAGE_INVALID_MODEL
            status = 7
        elif code == "50511":
            err_code = StatusCodes.OUTPUT_VIDEO_INVALID_MODEL
            status = 7
        else:
            err_code = ''
            status = 2
        video_result = {
            "width": 0,
            "height": 0,
            "fps": 0,
            "status": status,
            "tokens": 0,
            "save_path": None,
            "msg": msg,
            "code": err_code,
        }
        return video_result
    # 3. 下载并保存视频到本地
    # 生成带任务ID的文件名（便于追踪）
    filename = f"/video_task_{task_id}.mp4" if task_id else None
    download_video_result = download_single_video(
        video_url,
        save_dir + filename,
    )
    if not download_video_result:
        video_result["msg"] = "视频下载失败"
        return video_result
    video_info = get_video_info(save_dir + filename)
    print(f"视频保存成功：{save_dir + filename}")

    return {
        "tokens": tokens,
        'save_path': save_dir + filename,
        'width': video_info['width'],
        'height': video_info['height'],
        'duration': video_info['duration'],
        'fps': video_info['fps'],
        'code': "",
        'status': 1,
    }


if __name__ == '__main__':
    # task_params = {'taskId': 'f252c8484bc4f3367500', 'prompt': '跳舞', 'first_image': 'video/image_to_video/upload/d41d8cd98f00b204e9800998ecf8427e/1754404741567.jfif', 'last_image': '', 'duration': '5', 'fps': '30', 'resolution': '720p', 'ratio': '', 'appId': '1', 'taskType': 'pwai_video_itf', 'QueueResult': 'AILocalVideoResult', 'storageService': 'r2', 'isLogin': 1, 'bucketName': 'picwand', 'storedPrefix': 'video/image_to_video', 'progress': 1, 'queueName': 'AIHubVideoTasks_Test',
    #                'modelType': 2,
    #                'moduleType': 21, 'moduleName': 'kling-v1', 'mResolution': '720p', 'mode': 'std', 'negative_prompt': '', 'cfg_scale': 0}

    task_params = {
        'taskId': '222101afbb815d77d25e',
        'prompt': 'hug',
        'first_image': 'video/image_to_video/upload/d41d8cd98f00b204e9800998ecf8427e/1754989139319.jpg',
        'last_image': '',
        'duration': '5',
        'fps': '24',
        'resolution': '480p',
        'ratio': 'adaptive',
        'appId': '1',
        'taskType': 'pwai_video_itf',
        'QueueResult': 'AILocalVideoResult',
        'storageService': 'r2',
        'isLogin': 1,
        'bucketName': 'picwand',
        'storedPrefix': 'video/image_to_video',
        'progress': 1,
        'queueName': 'AIHubVideoTasks_Test',
        'modelType': 5,
        'moduleType': 20,
        'moduleName': 'doubao-seedance-1-0-pro',
        'mResolution': '480p',
        'mode': 'std',
        'negative_prompt': '',
        'cfg_scale': 0.5,
        'p_t_b': 0,
        'n_t_b': 0,
        'bgm': 0,

    }
    # task_params = {'taskId': '122521d7bb7c49145129', 'prompt': '希特勒殴打犹太人\n', 'first_image': '', 'last_image': '', 'duration': '5', 'fps': '24', 'resolution': '480p', 'ratio': '16:9', 'appId': '1', 'taskType': 'pwai_video_itf', 'QueueResult': 'AILocalVideoResult', 'storageService': 'r2', 'isLogin': 1, 'bucketName': 'picwand', 'storedPrefix': 'video/image_to_video', 'progress': 1, 'queueName': 'AIHubVideoTasks_Test', 'modelType': 4, 'moduleType': 20,
    #                'moduleName': 'doubao-seedance-1-0-pro', 'mResolution': '480p', 'mode': 'std', 'negative_prompt': '', 'cfg_scale': 0.5, 'p_t_b': 0, 'n_t_b': 0, 'motion_type': '', 'model_style': '', 'camera_movement': ''}
    # video_compress_image_size_path('./upload/img_1.png', 10, (300,6000), (300,6000))
    task_params = {'taskId': 'c2d58a51d6ad223b5334', 'prompt': 'Dance', 'first_image': 'video/image_to_video/upload/d41d8cd98f00b204e9800998ecf8427e/1759990457112.png', 'last_image': '', 'duration': '10', 'fps': '24', 'resolution': '480p', 'ratio': 'adaptive', 'appId': '101', 'taskType': 'pwai_video_itf', 'QueueResult': 'AIHubVideoResult', 'storageService': 'r2', 'isLogin': 1, 'bucketName': 'picwand', 'storedPrefix': 'video/image_to_video', 'progress': 1, 'queueName': 'AIHubVideoTasks', 'modelType': 3, 'moduleType': 20, 'moduleName': 'doubao-seedance-1-0-lite-i2v', 'mResolution': '480p', 'mode': 'std', 'negative_prompt': '', 'cfg_scale': 0.5, 'p_t_b': 0, 'n_t_b': 0, 'motion_type': '', 'model_style': '', 'camera_movement': '', 'template': "put_on_bunny_girl_outfit_720p"}
    # task_params =  {'taskId': '9dc080e575cad4197ac6', 'prompt': 'raise her legs but with fall', 'first_image': 'video/image_to_video/upload/d41d8cd98f00b204e9800998ecf8427e/1754639988027.jpeg', 'last_image': '', 'duration': '5', 'fps': '24', 'resolution': '480p', 'ratio': 'adaptive', 'appId': '1', 'taskType': 'pwai_video_itf', 'QueueResult': 'AILocalVideoResult', 'storageService': 'r2', 'isLogin': 1, 'bucketName': 'picwand', 'storedPrefix': 'video/image_to_video', 'progress': 1, 'queueName': 'AIHubVideoTasks_Test', 'modelType': 5, 'moduleType': 20, 'moduleName': 'doubao-seedance-1-0-pro', 'mResolution': '480p', 'mode': 'std', 'negative_prompt': '', 'cfg_scale': 0.5, 'p_t_b': 0, 'n_t_b': 0}
    video_generator = viduGenerator_doubao(task_params)
    print(video_generator)
